{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distributed data parallel MNIST training with PyTorch and SMDataParallel\n",
    "\n",
    "\n",
    "## Background\n",
    "SMDataParallel is a new capability in Amazon SageMaker to train deep learning models faster and cheaper. SMDataParallel is a distributed data parallel training framework for PyTorch. \n",
    "\n",
    "This notebook example shows how to use SMDataParallel with PyTorch in SageMaker using MNIST dataset.\n",
    "\n",
    "For more information:\n",
    "1. [PyTorch in SageMaker](https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html)\n",
    "2. [SMDataParallel PyTorch API Specification] < LINK TO BE ADDED >\n",
    "3. [Getting started with SMDataParallel on SageMaker] < LINK TO BE ADDED >\n",
    "\n",
    "**NOTE:** This example requires SageMaker Python SDK v2.X.\n",
    "\n",
    "\n",
    "### Dataset\n",
    "This example uses the MNIST dataset. MNIST is a widely used dataset for handwritten digit classification. It consists of 70,000 labeled 28x28 pixel grayscale images of hand-written digits. The dataset is split into 60,000 training images and 10,000 test images. There are 10 classes (one for each of the 10 digits).\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### SageMaker execution roles\n",
    "\n",
    "The IAM role arn used to give training and hosting access to your data. See the [Amazon SageMaker Roles](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-roles.html) for how to create these. Note, if more than one role is required for notebook instances, training, and/or hosting, please replace the sagemaker.get_execution_role() with the appropriate full IAM role arn string(s)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "\n",
    "sagemaker_session = sagemaker.Session()\n",
    "role = sagemaker.get_execution_role()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model training with SMDataParallel\n",
    "\n",
    "### Training script\n",
    "\n",
    "The MNIST dataset is downloaded using the `torchvision.datasets` PyTorch module; you can see how this is implemented in the `train_pytorch_smdataparallel_mnist.py` training script that is printed out in the next cell.\n",
    "\n",
    "The training script provides the code you need for distributed data parallel (DDP) training using SMDataParallel. The training script is very similar to a PyTorch training script you might run outside of SageMaker, but modified to run with SMDataParallel. SMDataParallel's PyTorch client provides an alternative to PyTorch's native DDP. For details about how to use SMDataParallel's DDP in your native PyTorch script, see the Getting Started with SMDataParallel tutorials.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[37m# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.\u001b[39;49;00m\r\n",
      "\u001b[37m#\u001b[39;49;00m\r\n",
      "\u001b[37m# Licensed under the Apache License, Version 2.0 (the \"License\"). You\u001b[39;49;00m\r\n",
      "\u001b[37m# may not use this file except in compliance with the License. A copy of\u001b[39;49;00m\r\n",
      "\u001b[37m# the License is located at\u001b[39;49;00m\r\n",
      "\u001b[37m#\u001b[39;49;00m\r\n",
      "\u001b[37m#     http://aws.amazon.com/apache2.0/\u001b[39;49;00m\r\n",
      "\u001b[37m#\u001b[39;49;00m\r\n",
      "\u001b[37m# or in the \"license\" file accompanying this file. This file is\u001b[39;49;00m\r\n",
      "\u001b[37m# distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF\u001b[39;49;00m\r\n",
      "\u001b[37m# ANY KIND, either express or implied. See the License for the specific\u001b[39;49;00m\r\n",
      "\u001b[37m# language governing permissions and limitations under the License.\u001b[39;49;00m\r\n",
      "\r\n",
      "\u001b[34mfrom\u001b[39;49;00m \u001b[04m\u001b[36m__future__\u001b[39;49;00m \u001b[34mimport\u001b[39;49;00m print_function\r\n",
      "\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mos\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36margparse\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mtime\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mtorch\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mtorch\u001b[39;49;00m\u001b[04m\u001b[36m.\u001b[39;49;00m\u001b[04m\u001b[36mnn\u001b[39;49;00m\u001b[04m\u001b[36m.\u001b[39;49;00m\u001b[04m\u001b[36mfunctional\u001b[39;49;00m \u001b[34mas\u001b[39;49;00m \u001b[04m\u001b[36mF\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mtorch\u001b[39;49;00m\u001b[04m\u001b[36m.\u001b[39;49;00m\u001b[04m\u001b[36moptim\u001b[39;49;00m \u001b[34mas\u001b[39;49;00m \u001b[04m\u001b[36moptim\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mtorch\u001b[39;49;00m\u001b[04m\u001b[36m.\u001b[39;49;00m\u001b[04m\u001b[36mnn\u001b[39;49;00m \u001b[34mas\u001b[39;49;00m \u001b[04m\u001b[36mnn\u001b[39;49;00m\r\n",
      "\u001b[34mfrom\u001b[39;49;00m \u001b[04m\u001b[36mtorchvision\u001b[39;49;00m \u001b[34mimport\u001b[39;49;00m datasets, transforms\r\n",
      "\u001b[34mfrom\u001b[39;49;00m \u001b[04m\u001b[36mtorch\u001b[39;49;00m\u001b[04m\u001b[36m.\u001b[39;49;00m\u001b[04m\u001b[36moptim\u001b[39;49;00m\u001b[04m\u001b[36m.\u001b[39;49;00m\u001b[04m\u001b[36mlr_scheduler\u001b[39;49;00m \u001b[34mimport\u001b[39;49;00m StepLR\r\n",
      "\r\n",
      "\u001b[37m# Network definition\u001b[39;49;00m\r\n",
      "\u001b[34mfrom\u001b[39;49;00m \u001b[04m\u001b[36mmodel_def\u001b[39;49;00m \u001b[34mimport\u001b[39;49;00m Net\r\n",
      "\r\n",
      "\u001b[37m# Import SMDataParallel PyTorch Modules\u001b[39;49;00m\r\n",
      "\u001b[34mfrom\u001b[39;49;00m \u001b[04m\u001b[36msmdistributed\u001b[39;49;00m\u001b[04m\u001b[36m.\u001b[39;49;00m\u001b[04m\u001b[36mdataparallel\u001b[39;49;00m\u001b[04m\u001b[36m.\u001b[39;49;00m\u001b[04m\u001b[36mtorch\u001b[39;49;00m\u001b[04m\u001b[36m.\u001b[39;49;00m\u001b[04m\u001b[36mparallel\u001b[39;49;00m\u001b[04m\u001b[36m.\u001b[39;49;00m\u001b[04m\u001b[36mdistributed\u001b[39;49;00m \u001b[34mimport\u001b[39;49;00m DistributedDataParallel \u001b[34mas\u001b[39;49;00m DDP\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36msmdistributed\u001b[39;49;00m\u001b[04m\u001b[36m.\u001b[39;49;00m\u001b[04m\u001b[36mdataparallel\u001b[39;49;00m\u001b[04m\u001b[36m.\u001b[39;49;00m\u001b[04m\u001b[36mtorch\u001b[39;49;00m\u001b[04m\u001b[36m.\u001b[39;49;00m\u001b[04m\u001b[36mdistributed\u001b[39;49;00m \u001b[34mas\u001b[39;49;00m \u001b[04m\u001b[36mdist\u001b[39;49;00m\r\n",
      "\r\n",
      "dist.init_process_group()\r\n",
      "\r\n",
      "\u001b[34mdef\u001b[39;49;00m \u001b[32mtrain\u001b[39;49;00m(args, model, device, train_loader, optimizer, epoch):\r\n",
      "    model.train()\r\n",
      "    \u001b[34mfor\u001b[39;49;00m batch_idx, (data, target) \u001b[35min\u001b[39;49;00m \u001b[36menumerate\u001b[39;49;00m(train_loader):\r\n",
      "        data, target = data.to(device), target.to(device)\r\n",
      "        optimizer.zero_grad()\r\n",
      "        output = model(data)\r\n",
      "        loss = F.nll_loss(output, target)\r\n",
      "        loss.backward()\r\n",
      "        optimizer.step()\r\n",
      "        \u001b[34mif\u001b[39;49;00m batch_idx % args.log_interval == \u001b[34m0\u001b[39;49;00m \u001b[35mand\u001b[39;49;00m args.rank == \u001b[34m0\u001b[39;49;00m:\r\n",
      "            \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mTrain Epoch: \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m [\u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m/\u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m (\u001b[39;49;00m\u001b[33m{:.0f}\u001b[39;49;00m\u001b[33m%\u001b[39;49;00m\u001b[33m)]\u001b[39;49;00m\u001b[33m\\t\u001b[39;49;00m\u001b[33mLoss: \u001b[39;49;00m\u001b[33m{:.6f}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(\r\n",
      "                epoch, batch_idx * \u001b[36mlen\u001b[39;49;00m(data) * args.world_size, \u001b[36mlen\u001b[39;49;00m(train_loader.dataset),\r\n",
      "                \u001b[34m100.\u001b[39;49;00m * batch_idx / \u001b[36mlen\u001b[39;49;00m(train_loader), loss.item()))\r\n",
      "        \u001b[34mif\u001b[39;49;00m args.verbose:\r\n",
      "            \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mBatch\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, batch_idx, \u001b[33m\"\u001b[39;49;00m\u001b[33mfrom rank\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m, args.rank)\r\n",
      "\r\n",
      "\r\n",
      "\u001b[34mdef\u001b[39;49;00m \u001b[32mtest\u001b[39;49;00m(model, device, test_loader):\r\n",
      "    model.eval()\r\n",
      "    test_loss = \u001b[34m0\u001b[39;49;00m\r\n",
      "    correct = \u001b[34m0\u001b[39;49;00m\r\n",
      "    \u001b[34mwith\u001b[39;49;00m torch.no_grad():\r\n",
      "        \u001b[34mfor\u001b[39;49;00m data, target \u001b[35min\u001b[39;49;00m test_loader:\r\n",
      "            data, target = data.to(device), target.to(device)\r\n",
      "            output = model(data)\r\n",
      "            test_loss += F.nll_loss(output, target, reduction=\u001b[33m'\u001b[39;49;00m\u001b[33msum\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m).item()  \u001b[37m# sum up batch loss\u001b[39;49;00m\r\n",
      "            pred = output.argmax(dim=\u001b[34m1\u001b[39;49;00m, keepdim=\u001b[34mTrue\u001b[39;49;00m)  \u001b[37m# get the index of the max log-probability\u001b[39;49;00m\r\n",
      "            correct += pred.eq(target.view_as(pred)).sum().item()\r\n",
      "\r\n",
      "    test_loss /= \u001b[36mlen\u001b[39;49;00m(test_loader.dataset)\r\n",
      "\r\n",
      "    \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33m\\n\u001b[39;49;00m\u001b[33mTest set: Average loss: \u001b[39;49;00m\u001b[33m{:.4f}\u001b[39;49;00m\u001b[33m, Accuracy: \u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m/\u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m (\u001b[39;49;00m\u001b[33m{:.0f}\u001b[39;49;00m\u001b[33m%\u001b[39;49;00m\u001b[33m)\u001b[39;49;00m\u001b[33m\\n\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(\r\n",
      "        test_loss, correct, \u001b[36mlen\u001b[39;49;00m(test_loader.dataset),\r\n",
      "        \u001b[34m100.\u001b[39;49;00m * correct / \u001b[36mlen\u001b[39;49;00m(test_loader.dataset)))\r\n",
      "\r\n",
      "\r\n",
      "\u001b[34mdef\u001b[39;49;00m \u001b[32msave_model\u001b[39;49;00m(model, model_dir):\r\n",
      "    \u001b[34mwith\u001b[39;49;00m \u001b[36mopen\u001b[39;49;00m(os.path.join(model_dir, \u001b[33m'\u001b[39;49;00m\u001b[33mmodel.pth\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m), \u001b[33m'\u001b[39;49;00m\u001b[33mwb\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m) \u001b[34mas\u001b[39;49;00m f:\r\n",
      "        torch.save(model.module.state_dict(), f)\r\n",
      "\r\n",
      "\u001b[34mdef\u001b[39;49;00m \u001b[32mmain\u001b[39;49;00m():\r\n",
      "    \u001b[37m# Training settings\u001b[39;49;00m\r\n",
      "    parser = argparse.ArgumentParser(description=\u001b[33m'\u001b[39;49;00m\u001b[33mPyTorch MNIST Example\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--batch-size\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mint\u001b[39;49;00m, default=\u001b[34m64\u001b[39;49;00m, metavar=\u001b[33m'\u001b[39;49;00m\u001b[33mN\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        help=\u001b[33m'\u001b[39;49;00m\u001b[33minput batch size for training (default: 64)\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--test-batch-size\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mint\u001b[39;49;00m, default=\u001b[34m1000\u001b[39;49;00m, metavar=\u001b[33m'\u001b[39;49;00m\u001b[33mN\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        help=\u001b[33m'\u001b[39;49;00m\u001b[33minput batch size for testing (default: 1000)\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--epochs\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mint\u001b[39;49;00m, default=\u001b[34m14\u001b[39;49;00m, metavar=\u001b[33m'\u001b[39;49;00m\u001b[33mN\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        help=\u001b[33m'\u001b[39;49;00m\u001b[33mnumber of epochs to train (default: 14)\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--lr\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mfloat\u001b[39;49;00m, default=\u001b[34m1.0\u001b[39;49;00m, metavar=\u001b[33m'\u001b[39;49;00m\u001b[33mLR\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        help=\u001b[33m'\u001b[39;49;00m\u001b[33mlearning rate (default: 1.0)\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--gamma\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mfloat\u001b[39;49;00m, default=\u001b[34m0.7\u001b[39;49;00m, metavar=\u001b[33m'\u001b[39;49;00m\u001b[33mM\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        help=\u001b[33m'\u001b[39;49;00m\u001b[33mLearning rate step gamma (default: 0.7)\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--seed\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mint\u001b[39;49;00m, default=\u001b[34m1\u001b[39;49;00m, metavar=\u001b[33m'\u001b[39;49;00m\u001b[33mS\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        help=\u001b[33m'\u001b[39;49;00m\u001b[33mrandom seed (default: 1)\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--log-interval\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mint\u001b[39;49;00m, default=\u001b[34m10\u001b[39;49;00m, metavar=\u001b[33m'\u001b[39;49;00m\u001b[33mN\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        help=\u001b[33m'\u001b[39;49;00m\u001b[33mhow many batches to wait before logging training status\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--save-model\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, action=\u001b[33m'\u001b[39;49;00m\u001b[33mstore_true\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, default=\u001b[34mFalse\u001b[39;49;00m,\r\n",
      "                        help=\u001b[33m'\u001b[39;49;00m\u001b[33mFor Saving the current Model\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--verbose\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, action=\u001b[33m'\u001b[39;49;00m\u001b[33mstore_true\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, default=\u001b[34mFalse\u001b[39;49;00m,\r\n",
      "                        help=\u001b[33m'\u001b[39;49;00m\u001b[33mFor displaying SMDataParallel-specific logs\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--data-path\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mstr\u001b[39;49;00m, default=\u001b[33m'\u001b[39;49;00m\u001b[33m/tmp/data\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, help=\u001b[33m'\u001b[39;49;00m\u001b[33mPath for downloading \u001b[39;49;00m\u001b[33m'\u001b[39;49;00m\r\n",
      "                                                                           \u001b[33m'\u001b[39;49;00m\u001b[33mthe MNIST dataset\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    \u001b[37m# Model checkpoint location\u001b[39;49;00m\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--model-dir\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mstr\u001b[39;49;00m, default=os.environ[\u001b[33m'\u001b[39;49;00m\u001b[33mSM_MODEL_DIR\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m])\r\n",
      "\r\n",
      "    args = parser.parse_args()\r\n",
      "    args.world_size = dist.get_world_size()\r\n",
      "    args.rank = rank = dist.get_rank()\r\n",
      "    args.local_rank = local_rank = dist.get_local_rank()\r\n",
      "    args.lr = \u001b[34m1.0\u001b[39;49;00m\r\n",
      "    args.batch_size //= args.world_size // \u001b[34m8\u001b[39;49;00m\r\n",
      "    args.batch_size = \u001b[36mmax\u001b[39;49;00m(args.batch_size, \u001b[34m1\u001b[39;49;00m)\r\n",
      "    data_path = args.data_path\r\n",
      "\r\n",
      "                        \r\n",
      "    \u001b[34mif\u001b[39;49;00m args.verbose:\r\n",
      "        \u001b[36mprint\u001b[39;49;00m(\u001b[33m'\u001b[39;49;00m\u001b[33mHello from rank\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, rank, \u001b[33m'\u001b[39;49;00m\u001b[33mof local_rank\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                local_rank, \u001b[33m'\u001b[39;49;00m\u001b[33min world size of\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, args.world_size)\r\n",
      "\r\n",
      "    \u001b[34mif\u001b[39;49;00m \u001b[35mnot\u001b[39;49;00m torch.cuda.is_available():\r\n",
      "        \u001b[34mraise\u001b[39;49;00m \u001b[36mException\u001b[39;49;00m(\u001b[33m\"\u001b[39;49;00m\u001b[33mMust run SMDataParallel MNIST example on CUDA-capable devices.\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m)\r\n",
      "\r\n",
      "    torch.manual_seed(args.seed)\r\n",
      "\r\n",
      "    device = torch.device(\u001b[33m\"\u001b[39;49;00m\u001b[33mcuda\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m)\r\n",
      "\r\n",
      "    \u001b[34mif\u001b[39;49;00m local_rank == \u001b[34m0\u001b[39;49;00m:\r\n",
      "        train_dataset = datasets.MNIST(data_path, train=\u001b[34mTrue\u001b[39;49;00m, download=\u001b[34mTrue\u001b[39;49;00m,\r\n",
      "                       transform=transforms.Compose([\r\n",
      "                           transforms.ToTensor(),\r\n",
      "                           transforms.Normalize((\u001b[34m0.1307\u001b[39;49;00m,), (\u001b[34m0.3081\u001b[39;49;00m,))\r\n",
      "                       ]))\r\n",
      "    \u001b[34melse\u001b[39;49;00m:\r\n",
      "        time.sleep(\u001b[34m8\u001b[39;49;00m)\r\n",
      "        train_dataset = datasets.MNIST(data_path, train=\u001b[34mTrue\u001b[39;49;00m, download=\u001b[34mFalse\u001b[39;49;00m,\r\n",
      "                       transform=transforms.Compose([\r\n",
      "                           transforms.ToTensor(),\r\n",
      "                           transforms.Normalize((\u001b[34m0.1307\u001b[39;49;00m,), (\u001b[34m0.3081\u001b[39;49;00m,))\r\n",
      "                       ]))\r\n",
      "\r\n",
      "    train_sampler = torch.utils.data.distributed.DistributedSampler(\r\n",
      "            train_dataset,\r\n",
      "            num_replicas=args.world_size,\r\n",
      "            rank=rank)\r\n",
      "    train_loader = torch.utils.data.DataLoader(\r\n",
      "        train_dataset,\r\n",
      "        batch_size=args.batch_size,\r\n",
      "        shuffle=\u001b[34mFalse\u001b[39;49;00m,\r\n",
      "        num_workers=\u001b[34m0\u001b[39;49;00m,\r\n",
      "        pin_memory=\u001b[34mTrue\u001b[39;49;00m,\r\n",
      "        sampler=train_sampler)\r\n",
      "    \u001b[34mif\u001b[39;49;00m rank == \u001b[34m0\u001b[39;49;00m:\r\n",
      "        test_loader = torch.utils.data.DataLoader(\r\n",
      "            datasets.MNIST(data_path, train=\u001b[34mFalse\u001b[39;49;00m, transform=transforms.Compose([\r\n",
      "                               transforms.ToTensor(),\r\n",
      "                               transforms.Normalize((\u001b[34m0.1307\u001b[39;49;00m,), (\u001b[34m0.3081\u001b[39;49;00m,))\r\n",
      "                           ])),\r\n",
      "            batch_size=args.test_batch_size, shuffle=\u001b[34mTrue\u001b[39;49;00m)\r\n",
      "    \r\n",
      "    \u001b[37m# Use SMDataParallel PyTorch DDP for efficient distributed training\u001b[39;49;00m\r\n",
      "    model = DDP(Net().to(device))\r\n",
      "    torch.cuda.set_device(local_rank)\r\n",
      "    model.cuda(local_rank)\r\n",
      "    optimizer = optim.Adadelta(model.parameters(), lr=args.lr)\r\n",
      "    scheduler = StepLR(optimizer, step_size=\u001b[34m1\u001b[39;49;00m, gamma=args.gamma)\r\n",
      "    \u001b[34mfor\u001b[39;49;00m epoch \u001b[35min\u001b[39;49;00m \u001b[36mrange\u001b[39;49;00m(\u001b[34m1\u001b[39;49;00m, args.epochs + \u001b[34m1\u001b[39;49;00m):\r\n",
      "        train(args, model, device, train_loader, optimizer, epoch)\r\n",
      "        \u001b[34mif\u001b[39;49;00m rank == \u001b[34m0\u001b[39;49;00m:\r\n",
      "            test(model, device, test_loader)\r\n",
      "        scheduler.step()\r\n",
      "\r\n",
      "    \u001b[34mif\u001b[39;49;00m rank == \u001b[34m0\u001b[39;49;00m:\r\n",
      "        \u001b[36mprint\u001b[39;49;00m(\u001b[33m\"\u001b[39;49;00m\u001b[33mSaving the model...\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m)\r\n",
      "        save_model(model, args.model_dir)   \r\n",
      "\r\n",
      "\r\n",
      "\u001b[34mif\u001b[39;49;00m \u001b[31m__name__\u001b[39;49;00m == \u001b[33m'\u001b[39;49;00m\u001b[33m__main__\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m:\r\n",
      "    main()\r\n"
     ]
    }
   ],
   "source": [
    "!pygmentize code/train_pytorch_smdataparallel_mnist.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Estimator function options\n",
    "\n",
    "In the following code block, you can update the estimator function to use a different instance type, instance count, and distrubtion strategy. You're also passing in the training script you reviewed in the previous cell.\n",
    "\n",
    "**Instance types**\n",
    "\n",
    "SMDataParallel supports model training on SageMaker with the following instance types only:\n",
    "1. ml.p3.16xlarge\n",
    "1. ml.p3dn.24xlarge [Recommended]\n",
    "1. ml.p4d.24xlarge [Recommended]\n",
    "\n",
    "**Instance count**\n",
    "\n",
    "To get the best performance and the most out of SMDataParallel, you should use at least 2 instances, but you can also use 1 for testing this example.\n",
    "\n",
    "**Distribution strategy**\n",
    "\n",
    "Note that to use DDP mode, you update the the `distribution` strategy, and set it to use `smdistributed dataparallel`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.pytorch import PyTorch\n",
    "estimator = PyTorch(base_job_name='pytorch-smdataparallel-mnist',\n",
    "                        source_dir='code',\n",
    "                        entry_point='train_pytorch_smdataparallel_mnist.py',\n",
    "                        role=role,\n",
    "                        framework_version='1.6.0',\n",
    "                        py_version='py3',\n",
    "                        # For training with multinode distributed training, set this count. Example: 2\n",
    "                        instance_count=2,\n",
    "                        # For training with p3dn instance use - ml.p3dn.24xlarge\n",
    "                        instance_type= 'ml.p3.16xlarge',\n",
    "                        sagemaker_session=sagemaker_session,\n",
    "                        # Training using SMDataParallel Distributed Training Framework\n",
    "                        distribution={'smdistributed':{\n",
    "                                            'dataparallel':{\n",
    "                                                    'enabled': True\n",
    "                                                 }\n",
    "                                          }\n",
    "                                      },\n",
    "                        debugger_hook_config=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2020-12-10 20:14:55 Starting - Starting the training job...\n",
      "2020-12-10 20:14:58 Starting - Launching requested ML instances.........\n",
      "2020-12-10 20:16:45 Starting - Preparing the instances for training.........\n",
      "2020-12-10 20:18:26 Downloading - Downloading input data\n",
      "2020-12-10 20:18:26 Training - Downloading the training image.........\n",
      "2020-12-10 20:19:56 Uploading - Uploading generated training model\u001b[34mbash: cannot set terminal process group (-1): Inappropriate ioctl for device\u001b[0m\n",
      "\u001b[34mbash: no job control in this shell\u001b[0m\n",
      "\u001b[34m2020-12-10 20:19:47,428 sagemaker-training-toolkit INFO     Imported framework sagemaker_pytorch_container.training\u001b[0m\n",
      "\u001b[34m2020-12-10 20:19:47,505 sagemaker_pytorch_container.training INFO     Block until all host DNS lookups succeed.\u001b[0m\n",
      "\u001b[34m2020-12-10 20:19:53,768 sagemaker_pytorch_container.training INFO     Invoking user training script.\u001b[0m\n",
      "\u001b[34mbash: cannot set terminal process group (-1): Inappropriate ioctl for device\u001b[0m\n",
      "\u001b[34mbash: no job control in this shell\u001b[0m\n",
      "\u001b[34m2020-12-10 20:19:51,055 sagemaker-training-toolkit INFO     Imported framework sagemaker_pytorch_container.training\u001b[0m\n",
      "\u001b[34m2020-12-10 20:19:51,133 sagemaker_pytorch_container.training INFO     Block until all host DNS lookups succeed.\u001b[0m\n",
      "\u001b[34m2020-12-10 20:19:52,604 sagemaker_pytorch_container.training INFO     Invoking user training script.\u001b[0m\n",
      "\u001b[34m2020-12-10 20:19:53,169 sagemaker-training-toolkit INFO     Invoking user script\n",
      "\u001b[0m\n",
      "\u001b[34mTraining Env:\n",
      "\u001b[0m\n",
      "\u001b[34m{\n",
      "    \"additional_framework_parameters\": {},\n",
      "    \"channel_input_dirs\": {},\n",
      "    \"current_host\": \"algo-2\",\n",
      "    \"framework_module\": \"sagemaker_pytorch_container.training:main\",\n",
      "    \"hosts\": [\n",
      "        \"algo-1\",\n",
      "        \"algo-2\"\n",
      "    ],\n",
      "    \"hyperparameters\": {},\n",
      "    \"input_config_dir\": \"/opt/ml/input/config\",\n",
      "    \"input_data_config\": {},\n",
      "    \"input_dir\": \"/opt/ml/input\",\n",
      "    \"is_master\": false,\n",
      "    \"job_name\": \"pytorch-smdataparallel-mnist-2020-12-10-20-14-54-846\",\n",
      "    \"log_level\": 20,\n",
      "    \"master_hostname\": \"algo-1\",\n",
      "    \"model_dir\": \"/opt/ml/model\",\n",
      "    \"module_dir\": \"s3://sagemaker-us-east-1-835319576252/pytorch-smdataparallel-mnist-2020-12-10-20-14-54-846/source/sourcedir.tar.gz\",\n",
      "    \"module_name\": \"train_pytorch_smdataparallel_mnist\",\n",
      "    \"network_interface_name\": \"eth0\",\n",
      "    \"num_cpus\": 64,\n",
      "    \"num_gpus\": 8,\n",
      "    \"output_data_dir\": \"/opt/ml/output/data\",\n",
      "    \"output_dir\": \"/opt/ml/output\",\n",
      "    \"output_intermediate_dir\": \"/opt/ml/output/intermediate\",\n",
      "    \"resource_config\": {\n",
      "        \"current_host\": \"algo-2\",\n",
      "        \"hosts\": [\n",
      "            \"algo-1\",\n",
      "            \"algo-2\"\n",
      "        ],\n",
      "        \"network_interface_name\": \"eth0\"\n",
      "    },\n",
      "    \"user_entry_point\": \"train_pytorch_smdataparallel_mnist.py\"\u001b[0m\n",
      "\u001b[34m}\n",
      "\u001b[0m\n",
      "\u001b[34mEnvironment variables:\n",
      "\u001b[0m\n",
      "\u001b[34mSM_HOSTS=[\"algo-1\",\"algo-2\"]\u001b[0m\n",
      "\u001b[34mSM_NETWORK_INTERFACE_NAME=eth0\u001b[0m\n",
      "\u001b[34mSM_HPS={}\u001b[0m\n",
      "\u001b[34mSM_USER_ENTRY_POINT=train_pytorch_smdataparallel_mnist.py\u001b[0m\n",
      "\u001b[34mSM_FRAMEWORK_PARAMS={}\u001b[0m\n",
      "\u001b[34mSM_RESOURCE_CONFIG={\"current_host\":\"algo-2\",\"hosts\":[\"algo-1\",\"algo-2\"],\"network_interface_name\":\"eth0\"}\u001b[0m\n",
      "\u001b[34mSM_INPUT_DATA_CONFIG={}\u001b[0m\n",
      "\u001b[34mSM_OUTPUT_DATA_DIR=/opt/ml/output/data\u001b[0m\n",
      "\u001b[34mSM_CHANNELS=[]\u001b[0m\n",
      "\u001b[34mSM_CURRENT_HOST=algo-2\u001b[0m\n",
      "\u001b[34mSM_MODULE_NAME=train_pytorch_smdataparallel_mnist\u001b[0m\n",
      "\u001b[34mSM_LOG_LEVEL=20\u001b[0m\n",
      "\u001b[34mSM_FRAMEWORK_MODULE=sagemaker_pytorch_container.training:main\u001b[0m\n",
      "\u001b[34mSM_INPUT_DIR=/opt/ml/input\u001b[0m\n",
      "\u001b[34mSM_INPUT_CONFIG_DIR=/opt/ml/input/config\u001b[0m\n",
      "\u001b[34mSM_OUTPUT_DIR=/opt/ml/output\u001b[0m\n",
      "\u001b[34mSM_NUM_CPUS=64\u001b[0m\n",
      "\u001b[34mSM_NUM_GPUS=8\u001b[0m\n",
      "\u001b[34mSM_MODEL_DIR=/opt/ml/model\u001b[0m\n",
      "\u001b[34mSM_MODULE_DIR=s3://sagemaker-us-east-1-835319576252/pytorch-smdataparallel-mnist-2020-12-10-20-14-54-846/source/sourcedir.tar.gz\u001b[0m\n",
      "\u001b[34mSM_TRAINING_ENV={\"additional_framework_parameters\":{},\"channel_input_dirs\":{},\"current_host\":\"algo-2\",\"framework_module\":\"sagemaker_pytorch_container.training:main\",\"hosts\":[\"algo-1\",\"algo-2\"],\"hyperparameters\":{},\"input_config_dir\":\"/opt/ml/input/config\",\"input_data_config\":{},\"input_dir\":\"/opt/ml/input\",\"is_master\":false,\"job_name\":\"pytorch-smdataparallel-mnist-2020-12-10-20-14-54-846\",\"log_level\":20,\"master_hostname\":\"algo-1\",\"model_dir\":\"/opt/ml/model\",\"module_dir\":\"s3://sagemaker-us-east-1-835319576252/pytorch-smdataparallel-mnist-2020-12-10-20-14-54-846/source/sourcedir.tar.gz\",\"module_name\":\"train_pytorch_smdataparallel_mnist\",\"network_interface_name\":\"eth0\",\"num_cpus\":64,\"num_gpus\":8,\"output_data_dir\":\"/opt/ml/output/data\",\"output_dir\":\"/opt/ml/output\",\"output_intermediate_dir\":\"/opt/ml/output/intermediate\",\"resource_config\":{\"current_host\":\"algo-2\",\"hosts\":[\"algo-1\",\"algo-2\"],\"network_interface_name\":\"eth0\"},\"user_entry_point\":\"train_pytorch_smdataparallel_mnist.py\"}\u001b[0m\n",
      "\u001b[34mSM_USER_ARGS=[]\u001b[0m\n",
      "\u001b[34mSM_OUTPUT_INTERMEDIATE_DIR=/opt/ml/output/intermediate\u001b[0m\n",
      "\u001b[34mPYTHONPATH=/opt/ml/code:/opt/conda/bin:/opt/conda/lib/python36.zip:/opt/conda/lib/python3.6:/opt/conda/lib/python3.6/lib-dynload:/opt/conda/lib/python3.6/site-packages\n",
      "\u001b[0m\n",
      "\u001b[34mInvoking script with the following command:\n",
      "\u001b[0m\n",
      "\u001b[34m/opt/conda/bin/python train_pytorch_smdataparallel_mnist.py\n",
      "\n",
      "\u001b[0m\n",
      "\u001b[35m2020-12-10 20:19:54,456 sagemaker-training-toolkit INFO     Invoking user script\n",
      "\u001b[0m\n",
      "\u001b[35mTraining Env:\n",
      "\u001b[0m\n",
      "\u001b[35m{\n",
      "    \"additional_framework_parameters\": {},\n",
      "    \"channel_input_dirs\": {},\n",
      "    \"current_host\": \"algo-1\",\n",
      "    \"framework_module\": \"sagemaker_pytorch_container.training:main\",\n",
      "    \"hosts\": [\n",
      "        \"algo-1\",\n",
      "        \"algo-2\"\n",
      "    ],\n",
      "    \"hyperparameters\": {},\n",
      "    \"input_config_dir\": \"/opt/ml/input/config\",\n",
      "    \"input_data_config\": {},\n",
      "    \"input_dir\": \"/opt/ml/input\",\n",
      "    \"is_master\": true,\n",
      "    \"job_name\": \"pytorch-smdataparallel-mnist-2020-12-10-20-14-54-846\",\n",
      "    \"log_level\": 20,\n",
      "    \"master_hostname\": \"algo-1\",\n",
      "    \"model_dir\": \"/opt/ml/model\",\n",
      "    \"module_dir\": \"s3://sagemaker-us-east-1-835319576252/pytorch-smdataparallel-mnist-2020-12-10-20-14-54-846/source/sourcedir.tar.gz\",\n",
      "    \"module_name\": \"train_pytorch_smdataparallel_mnist\",\n",
      "    \"network_interface_name\": \"eth0\",\n",
      "    \"num_cpus\": 64,\n",
      "    \"num_gpus\": 8,\n",
      "    \"output_data_dir\": \"/opt/ml/output/data\",\n",
      "    \"output_dir\": \"/opt/ml/output\",\n",
      "    \"output_intermediate_dir\": \"/opt/ml/output/intermediate\",\n",
      "    \"resource_config\": {\n",
      "        \"current_host\": \"algo-1\",\n",
      "        \"hosts\": [\n",
      "            \"algo-1\",\n",
      "            \"algo-2\"\n",
      "        ],\n",
      "        \"network_interface_name\": \"eth0\"\n",
      "    },\n",
      "    \"user_entry_point\": \"train_pytorch_smdataparallel_mnist.py\"\u001b[0m\n",
      "\u001b[35m}\n",
      "\u001b[0m\n",
      "\u001b[35mEnvironment variables:\n",
      "\u001b[0m\n",
      "\u001b[35mSM_HOSTS=[\"algo-1\",\"algo-2\"]\u001b[0m\n",
      "\u001b[35mSM_NETWORK_INTERFACE_NAME=eth0\u001b[0m\n",
      "\u001b[35mSM_HPS={}\u001b[0m\n",
      "\u001b[35mSM_USER_ENTRY_POINT=train_pytorch_smdataparallel_mnist.py\u001b[0m\n",
      "\u001b[35mSM_FRAMEWORK_PARAMS={}\u001b[0m\n",
      "\u001b[35mSM_RESOURCE_CONFIG={\"current_host\":\"algo-1\",\"hosts\":[\"algo-1\",\"algo-2\"],\"network_interface_name\":\"eth0\"}\u001b[0m\n",
      "\u001b[35mSM_INPUT_DATA_CONFIG={}\u001b[0m\n",
      "\u001b[35mSM_OUTPUT_DATA_DIR=/opt/ml/output/data\u001b[0m\n",
      "\u001b[35mSM_CHANNELS=[]\u001b[0m\n",
      "\u001b[35mSM_CURRENT_HOST=algo-1\u001b[0m\n",
      "\u001b[35mSM_MODULE_NAME=train_pytorch_smdataparallel_mnist\u001b[0m\n",
      "\u001b[35mSM_LOG_LEVEL=20\u001b[0m\n",
      "\u001b[35mSM_FRAMEWORK_MODULE=sagemaker_pytorch_container.training:main\u001b[0m\n",
      "\u001b[35mSM_INPUT_DIR=/opt/ml/input\u001b[0m\n",
      "\u001b[35mSM_INPUT_CONFIG_DIR=/opt/ml/input/config\u001b[0m\n",
      "\u001b[35mSM_OUTPUT_DIR=/opt/ml/output\u001b[0m\n",
      "\u001b[35mSM_NUM_CPUS=64\u001b[0m\n",
      "\u001b[35mSM_NUM_GPUS=8\u001b[0m\n",
      "\u001b[35mSM_MODEL_DIR=/opt/ml/model\u001b[0m\n",
      "\u001b[35mSM_MODULE_DIR=s3://sagemaker-us-east-1-835319576252/pytorch-smdataparallel-mnist-2020-12-10-20-14-54-846/source/sourcedir.tar.gz\u001b[0m\n",
      "\u001b[35mSM_TRAINING_ENV={\"additional_framework_parameters\":{},\"channel_input_dirs\":{},\"current_host\":\"algo-1\",\"framework_module\":\"sagemaker_pytorch_container.training:main\",\"hosts\":[\"algo-1\",\"algo-2\"],\"hyperparameters\":{},\"input_config_dir\":\"/opt/ml/input/config\",\"input_data_config\":{},\"input_dir\":\"/opt/ml/input\",\"is_master\":true,\"job_name\":\"pytorch-smdataparallel-mnist-2020-12-10-20-14-54-846\",\"log_level\":20,\"master_hostname\":\"algo-1\",\"model_dir\":\"/opt/ml/model\",\"module_dir\":\"s3://sagemaker-us-east-1-835319576252/pytorch-smdataparallel-mnist-2020-12-10-20-14-54-846/source/sourcedir.tar.gz\",\"module_name\":\"train_pytorch_smdataparallel_mnist\",\"network_interface_name\":\"eth0\",\"num_cpus\":64,\"num_gpus\":8,\"output_data_dir\":\"/opt/ml/output/data\",\"output_dir\":\"/opt/ml/output\",\"output_intermediate_dir\":\"/opt/ml/output/intermediate\",\"resource_config\":{\"current_host\":\"algo-1\",\"hosts\":[\"algo-1\",\"algo-2\"],\"network_interface_name\":\"eth0\"},\"user_entry_point\":\"train_pytorch_smdataparallel_mnist.py\"}\u001b[0m\n",
      "\u001b[35mSM_USER_ARGS=[]\u001b[0m\n",
      "\u001b[35mSM_OUTPUT_INTERMEDIATE_DIR=/opt/ml/output/intermediate\u001b[0m\n",
      "\u001b[35mPYTHONPATH=/opt/ml/code:/opt/conda/bin:/opt/conda/lib/python36.zip:/opt/conda/lib/python3.6:/opt/conda/lib/python3.6/lib-dynload:/opt/conda/lib/python3.6/site-packages\n",
      "\u001b[0m\n",
      "\u001b[35mInvoking script with the following command:\n",
      "\u001b[0m\n",
      "\u001b[35m/opt/conda/bin/python train_pytorch_smdataparallel_mnist.py\n",
      "\n",
      "\u001b[0m\n",
      "\u001b[34m2020-12-10 20:19:55,718 sagemaker-training-toolkit ERROR    ExecuteUserScriptError:\u001b[0m\n",
      "\u001b[34mCommand \"/opt/conda/bin/python train_pytorch_smdataparallel_mnist.py\"\u001b[0m\n",
      "\u001b[34mTraceback (most recent call last):\n",
      "  File \"train_pytorch_smdataparallel_mnist.py\", line 30, in <module>\n",
      "    from smdistributed.dataparallel.torch.parallel.distributed import DistributedDataParallel as DDP\u001b[0m\n",
      "\u001b[34mModuleNotFoundError: No module named 'smdistributed'\u001b[0m\n",
      "\u001b[35m2020-12-10 20:19:57,655 sagemaker-training-toolkit ERROR    ExecuteUserScriptError:\u001b[0m\n",
      "\u001b[35mCommand \"/opt/conda/bin/python train_pytorch_smdataparallel_mnist.py\"\u001b[0m\n",
      "\u001b[35mTraceback (most recent call last):\n",
      "  File \"train_pytorch_smdataparallel_mnist.py\", line 30, in <module>\n",
      "    from smdistributed.dataparallel.torch.parallel.distributed import DistributedDataParallel as DDP\u001b[0m\n",
      "\u001b[35mModuleNotFoundError: No module named 'smdistributed'\u001b[0m\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "2020-12-10 20:20:04 Failed - Training job failed\n"
     ]
    },
    {
     "ename": "UnexpectedStatusException",
     "evalue": "Error for Training job pytorch-smdataparallel-mnist-2020-12-10-20-14-54-846: Failed. Reason: AlgorithmError: ExecuteUserScriptError:\nCommand \"/opt/conda/bin/python train_pytorch_smdataparallel_mnist.py\"\nTraceback (most recent call last):\n  File \"train_pytorch_smdataparallel_mnist.py\", line 30, in <module>\n    from smdistributed.dataparallel.torch.parallel.distributed import DistributedDataParallel as DDP\nModuleNotFoundError: No module named 'smdistributed'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mUnexpectedStatusException\u001b[0m                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-4-28be9b2c12b5>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mestimator\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages/sagemaker/estimator.py\u001b[0m in \u001b[0;36mfit\u001b[0;34m(self, inputs, wait, logs, job_name, experiment_config)\u001b[0m\n\u001b[1;32m    533\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjobs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlatest_training_job\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    534\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mwait\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 535\u001b[0;31m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlatest_training_job\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwait\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlogs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mlogs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    536\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    537\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_compilation_job_name\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages/sagemaker/estimator.py\u001b[0m in \u001b[0;36mwait\u001b[0;34m(self, logs)\u001b[0m\n\u001b[1;32m   1208\u001b[0m         \u001b[0;31m# If logs are requested, call logs_for_jobs.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1209\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mlogs\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0;34m\"None\"\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1210\u001b[0;31m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msagemaker_session\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlogs_for_job\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjob_name\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mwait\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlog_type\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mlogs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1211\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1212\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msagemaker_session\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwait_for_job\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjob_name\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages/sagemaker/session.py\u001b[0m in \u001b[0;36mlogs_for_job\u001b[0;34m(self, job_name, wait, poll, log_type)\u001b[0m\n\u001b[1;32m   3363\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3364\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mwait\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3365\u001b[0;31m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_check_job_status\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mjob_name\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdescription\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"TrainingJobStatus\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   3366\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mdot\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3367\u001b[0m                 \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages/sagemaker/session.py\u001b[0m in \u001b[0;36m_check_job_status\u001b[0;34m(self, job, desc, status_key_name)\u001b[0m\n\u001b[1;32m   2955\u001b[0m                 ),\n\u001b[1;32m   2956\u001b[0m                 \u001b[0mallowed_statuses\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"Completed\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"Stopped\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2957\u001b[0;31m                 \u001b[0mactual_status\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mstatus\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2958\u001b[0m             )\n\u001b[1;32m   2959\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mUnexpectedStatusException\u001b[0m: Error for Training job pytorch-smdataparallel-mnist-2020-12-10-20-14-54-846: Failed. Reason: AlgorithmError: ExecuteUserScriptError:\nCommand \"/opt/conda/bin/python train_pytorch_smdataparallel_mnist.py\"\nTraceback (most recent call last):\n  File \"train_pytorch_smdataparallel_mnist.py\", line 30, in <module>\n    from smdistributed.dataparallel.torch.parallel.distributed import DistributedDataParallel as DDP\nModuleNotFoundError: No module named 'smdistributed'"
     ]
    }
   ],
   "source": [
    "estimator.fit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Next steps\n",
    "\n",
    "Now that you have a trained model, you can deploy an endpoint to host the model. After you deploy the endpoint, you can then test it with inference requests. The following cell will store the model_data variable to be used with the inference notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Storing s3://sagemaker-us-east-1-835319576252/pytorch-smdataparallel-mnist-2020-12-10-20-14-54-846/output/model.tar.gz as model_data\n",
      "Stored 'model_data' (str)\n"
     ]
    }
   ],
   "source": [
    "model_data = estimator.model_data\n",
    "print(\"Storing {} as model_data\".format(model_data))\n",
    "%store model_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_pytorch_p36",
   "language": "python",
   "name": "conda_pytorch_p36"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
